#!/bin/bash

# Jetson Nano Swap Optimization Script for 4GB Memory Limitation
# 专门针对4GB内存不足和递归代码卡死问题的swap优化脚本

echo "=== Jetson Nano Swap Optimization Script ==="
echo "针对4GB内存限制进行swap优化..."

# 检查当前内存和swap状态
echo "=== 当前内存状态 ==="
free -h
echo ""
echo "=== 当前swap状态 ==="
swapon --show
echo ""

# 备份当前配置
echo "备份当前系统配置..."
sudo cp /etc/fstab /etc/fstab.backup.$(date +%Y%m%d_%H%M%S)
sudo cp /etc/sysctl.conf /etc/sysctl.conf.backup.$(date +%Y%m%d_%H%M%S)

# 1. 创建或扩展swap文件
SWAP_SIZE="4G"  # 推荐swap大小为物理内存的1倍
SWAP_FILE="/swapfile"

echo "创建${SWAP_SIZE}的swap文件..."

# 检查是否已存在swap文件
if [ -f "$SWAP_FILE" ]; then
    echo "检测到现有swap文件，先禁用..."
    sudo swapoff "$SWAP_FILE"
    sudo rm "$SWAP_FILE"
fi

# 创建新的swap文件
echo "创建新的swap文件 (${SWAP_SIZE})..."
sudo fallocate -l "$SWAP_SIZE" "$SWAP_FILE"

# 如果fallocate失败，使用dd命令
if [ $? -ne 0 ]; then
    echo "fallocate失败，使用dd命令创建swap文件..."
    sudo dd if=/dev/zero of="$SWAP_FILE" bs=1M count=4096 status=progress
fi

# 设置正确的权限
sudo chmod 600 "$SWAP_FILE"

# 设置swap文件
echo "设置swap文件..."
sudo mkswap "$SWAP_FILE"

# 启用swap
echo "启用swap文件..."
sudo swapon "$SWAP_FILE"

# 2. 优化swap参数以防止递归代码卡死
echo "优化swap参数..."

# 创建swap优化配置
sudo tee /etc/sysctl.d/99-swap-optimization.conf > /dev/null << 'EOF'
# Swap optimization for 4GB memory limitation and recursive code protection

# 设置swappiness为60，在内存压力大时更积极使用swap
# 这有助于防止递归代码导致的内存耗尽
vm.swappiness = 60

# 优化内存回收策略
vm.vfs_cache_pressure = 100

# 设置脏页参数，防止大量脏页占用内存
vm.dirty_ratio = 10
vm.dirty_background_ratio = 5

# 设置最小空闲内存，保证系统响应性
vm.min_free_kbytes = 131072

# 优化内存分配策略
vm.overcommit_memory = 1
vm.overcommit_ratio = 50

# 防止内存碎片
vm.compact_memory = 1

# OOM killer优化，优先杀死占用内存最多的进程
vm.oom_kill_allocating_task = 1
vm.panic_on_oom = 0

# 网络缓冲区优化
net.core.rmem_max = 67108864
net.core.wmem_max = 67108864
net.core.netdev_max_backlog = 2500
EOF

# 应用swap优化参数
echo "应用swap优化参数..."
sudo sysctl -p /etc/sysctl.d/99-swap-optimization.conf

# 3. 添加swap到fstab以便开机自动挂载
echo "配置开机自动挂载swap..."
if ! grep -q "$SWAP_FILE" /etc/fstab; then
    echo "$SWAP_FILE none swap sw 0 0" | sudo tee -a /etc/fstab
    echo "已添加swap到fstab"
else
    echo "swap已在fstab中配置"
fi

# 4. 创建内存监控和保护脚本
echo "创建内存监控和保护脚本..."
sudo tee /usr/local/bin/memory-guard.sh > /dev/null << 'EOF'
#!/bin/bash
# 内存保护脚本，防止递归代码导致系统卡死

MEMORY_THRESHOLD=90  # 内存使用率阈值
SWAP_THRESHOLD=80    # swap使用率阈值

while true; do
    # 获取内存使用率
    MEMORY_USAGE=$(free | grep Mem | awk '{printf("%.0f", $3/$2 * 100.0)}')
    SWAP_USAGE=$(free | grep Swap | awk '{if($2>0) printf("%.0f", $3/$2 * 100.0); else print "0"}')
    
    # 检查内存使用率
    if [ "$MEMORY_USAGE" -gt "$MEMORY_THRESHOLD" ]; then
        echo "$(date): 警告 - 内存使用率过高: ${MEMORY_USAGE}%" | sudo tee -a /var/log/memory-guard.log
        
        # 强制回收缓存
        echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null
        
        # 如果swap使用率也很高，记录占用内存最多的进程
        if [ "$SWAP_USAGE" -gt "$SWAP_THRESHOLD" ]; then
            echo "$(date): 严重警告 - Swap使用率: ${SWAP_USAGE}%" | sudo tee -a /var/log/memory-guard.log
            echo "$(date): 占用内存最多的进程:" | sudo tee -a /var/log/memory-guard.log
            ps aux --sort=-%mem | head -5 | sudo tee -a /var/log/memory-guard.log
        fi
    fi
    
    sleep 30  # 每30秒检查一次
done
EOF

sudo chmod +x /usr/local/bin/memory-guard.sh

# 5. 创建systemd服务以自动运行内存保护
echo "创建内存保护服务..."
sudo tee /etc/systemd/system/memory-guard.service > /dev/null << 'EOF'
[Unit]
Description=Memory Guard Service
After=multi-user.target

[Service]
Type=simple
ExecStart=/usr/local/bin/memory-guard.sh
Restart=always
RestartSec=10
User=root

[Install]
WantedBy=multi-user.target
EOF

# 启用并启动内存保护服务
sudo systemctl daemon-reload
sudo systemctl enable memory-guard.service
sudo systemctl start memory-guard.service

# 6. 创建紧急内存清理脚本
echo "创建紧急内存清理脚本..."
sudo tee /usr/local/bin/emergency-memory-cleanup.sh > /dev/null << 'EOF'
#!/bin/bash
# 紧急内存清理脚本

echo "=== 紧急内存清理 $(date) ==="

# 1. 强制回收所有缓存
echo "清理系统缓存..."
sync
echo 3 > /proc/sys/vm/drop_caches

# 2. 清理临时文件
echo "清理临时文件..."
find /tmp -type f -atime +1 -delete 2>/dev/null
find /var/tmp -type f -atime +1 -delete 2>/dev/null

# 3. 清理日志文件
echo "清理旧日志文件..."
journalctl --vacuum-time=1d 2>/dev/null

# 4. 显示清理后的内存状态
echo "清理后内存状态："
free -h

echo "紧急清理完成"
EOF

sudo chmod +x /usr/local/bin/emergency-memory-cleanup.sh

# 7. 显示优化结果
echo ""
echo "=== Swap优化完成 ==="
echo "优化后内存和swap状态："
free -h
echo ""
echo "Swap文件信息："
swapon --show
echo ""
echo "已创建的工具："
echo "✅ ${SWAP_SIZE} swap文件: $SWAP_FILE"
echo "✅ 内存监控服务: memory-guard.service"
echo "✅ 紧急清理脚本: emergency-memory-cleanup.sh"
echo ""
echo "优化参数："
echo "• swappiness = 60 (更积极使用swap)"
echo "• min_free_kbytes = 131072 (保证128MB空闲内存)"
echo "• 内存使用率监控阈值: 90%"
echo "• Swap使用率监控阈值: 80%"
echo ""
echo "使用建议："
echo "1. 重启系统以确保所有优化生效"
echo "2. 运行 'sudo systemctl status memory-guard' 检查监控服务"
echo "3. 紧急情况下运行 'sudo emergency-memory-cleanup.sh'"
echo "4. 查看内存日志: 'sudo tail -f /var/log/memory-guard.log'"
echo ""
echo "针对递归代码的建议："
echo "• 在代码中添加递归深度限制"
echo "• 使用迭代替代深度递归"
echo "• 添加内存使用监控"
echo ""
echo "如需回滚swap设置："
echo "sudo swapoff $SWAP_FILE"
echo "sudo rm $SWAP_FILE"
echo "sudo cp /etc/fstab.backup.* /etc/fstab"